import json
import os.path
import shutil
import sqlite3
import tempfile

from fastapi.responses import FileResponse
from pydantic import BaseModel, Field

from datasets import Dataset


# 定义请求和响应数据的 Pydantic 模型
class FormData(BaseModel):
    question: list[str]
    ground_truths: list[str] = Field(alias="groundTruths")
    answer: list[str]
    contexts: list[list[str]]


class Database:
    def __init__(self, db_name='dataset.db'):
        self.conn = sqlite3.connect(db_name, check_same_thread=False)
        self.cursor = self.conn.cursor()
        self.create_table()

    def create_table(self):
        # 创建表格
        self.cursor.execute('''
        CREATE TABLE IF NOT EXISTS datasets (
            pid INTEGER PRIMARY KEY AUTOINCREMENT,
            data TEXT
        )
        ''')
        self.conn.commit()

    def insert_forms(self, data: dict):
        # 将字典数据转换为 JSON 字符串
        json_data = json.dumps(data)

        # 向表格中插入数据
        sql = """
            INSERT INTO datasets (data)
            VALUES (?)
        """
        self.cursor.execute(sql, (json_data,))
        self.conn.commit()

    def delete_dataset(self, pid: str):
        # 删除数据集
        sql = f"""
        DELETE FROM datasets WHERE pid = {pid}
        """
        self.cursor.execute(sql)
        self.conn.commit()

    def query_dataset(self):
        # 查询数据集
        sql = """
        SELECT pid, data FROM datasets ORDER BY pid DESC
        """
        self.cursor.execute(sql)
        results = self.cursor.fetchall()

        # 将每个结果转换为 Python 字典并存储在列表中
        result_list = []
        for row in results:
            data_dict = {"pid": row[0], "data": json.loads(row[1])}
            result_list.append(data_dict)  # 将 Python 字典添加到结果列表中

        return result_list

    def export_dataset(self, pid: str):
        # 查询数据集
        sql = f"""
        SELECT data FROM datasets WHERE pid = {pid}
        """
        self.cursor.execute(sql)
        results = self.cursor.fetchall()[0][0]

        # 将字符串 result 解析为字典对象
        result_dict = json.loads(results)
        # 使用字典创建 Dataset 对象
        dataset = Dataset.from_dict(result_dict)

        # 创建临时目录
        temp_dir = tempfile.mkdtemp()

        file_path = os.path.join(temp_dir, f"data_{pid}")
        print(file_path)
        # 保存数据集到临时目录
        dataset.save_to_disk(file_path)

        # 创建压缩文件
        zip_file_path = shutil.make_archive(base_name=file_path, root_dir=file_path, format='zip')
        print(zip_file_path)

        # 返回压缩文件给浏览器端
        result_zip = FileResponse(zip_file_path, media_type="application/zip")

        return result_zip
